import os
import re
import sys
import json
import math
import requests
import logging
import re

from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection

# -------------------- 配置部分 --------------------
DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")

MAX_SYMBOLS = 10  # 每个文件最大保留符号数
MIN_STR_LEN = 8  # 最小字符串长度

NAPI_PATTERN = r'^napi_[a-zA-Z0-9_]+_[a-zA-Z0-9_]+$'
STD_LIB_FILTER = {
    '__stack_chk', '__libc_', 'memset', 'memcpy',
    'printf', 'malloc', 'free', '_Z'  # C++符号
}

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers= [
        logging.StreamHandler(),
        logging.FileHandler('origin_analysis.log')
    ]
)

logger = logging.getLogger(__name__)

# -------------------- 工具函数 --------------------
def name_entropy(name):
    """计算符号名称熵值"""
    freq = {}
    for char in name:
        freq[char] = freq.get(char, 0) + 1
    entropy = 0.0
    for count in freq.values():
        p = count / len(name)
        entropy -= p * math.log2(p)
    return entropy


def is_obfuscated(name):
    """检测混淆符号"""
    return name_entropy(name) > 3.5 and re.match(r'^[a-f0-9]{16}$', name)


# -------------------- 核心分析类 --------------------
class SOAnalyzer:
    def __init__(self):
        self.symbol_rules = {
            'JNI': [NAPI_PATTERN]
        }

    @staticmethod
    def _filter_symbol(symbol):
        """基础符号过滤"""
        name = symbol.name
        if not name:
            return False
        if any(kw in name for kw in STD_LIB_FILTER):
            return False
        if symbol['st_size'] < 0x20:  # 过滤微小函数
            return False
        return symbol['st_info']['type'] == 'STT_FUNC'

    @staticmethod
    def _extract_strings(data):
        """改进的字符串提取"""
        strings = []
        current = []
        for byte in data:
            if 0x20 <= byte <= 0x7E:  # 可打印ASCII
                current.append(chr(byte))
            else:
                if len(current) >= MIN_STR_LEN:
                    strings.append(''.join(current))
                current = []
        return list(set(strings))  # 去重

    def analyze_elf(self, file_path):
        """执行完整分析"""
        result = {
            'file': os.path.basename(file_path),
            'symbols': [],
            'strings': [],
            'metadata': {}
        }

        with open(file_path, 'rb') as f:
            elf = ELFFile(f)

            # 提取元数据
            result['metadata'].update({
                'arch': elf.get_machine_arch(),
                'endian': elf.little_endian,
                'entry_point': hex(elf.header['e_entry'])
            })

            # 符号分析
            symbol_tables = [
                section for section in elf.iter_sections()
                if isinstance(section, SymbolTableSection)
            ]

            for section in symbol_tables:
                for symbol in section.iter_symbols():
                    if not self._filter_symbol(symbol):
                        continue

                    sym_info = {
                        'name': symbol.name,
                        'value': hex(symbol['st_value']),
                        'size': symbol['st_size'],
                        'section': elf.get_section(symbol['st_shndx']).name,
                        'entropy': round(name_entropy(symbol.name), 3),
                        'obfuscated': is_obfuscated(symbol.name)
                    }
                    result['symbols'].append(sym_info)

            # 字符串提取（优先.rodata和.dynstr）
            for section in elf.iter_sections():
                if section.name in ('.rodata', '.dynstr'):
                    result['strings'] += self._extract_strings(section.data())

            result['strings'] = list(filter(lambda x: any(kw in x.lower() for kw in ['author', 'copyright', 'version', '/', 'com.', 'http']), result['strings']))[:20]

            # 按优先级排序
            result['symbols'].sort(
                key=lambda x: (-x['size'], -x['entropy'])
            )
            result['symbols'] = list(map(lambda v: v['name'] , result['symbols'][:MAX_SYMBOLS]))

        return result


# -------------------- 情报查询类 --------------------
class ThreatIntelligence:
    @staticmethod
    def query_deepseek(context):
        combined_text = json.dumps(context)

        """调用DeepSeek API进行分析"""
        headers = {
            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": "You are a binary analysis expert"},
                {
                    "role": "user",
                    "content": f"""Analyze this library and determine its origin based on the following binary characteristics：{combined_text}
                        Please classify this libray into one of these broad categories:
                        1. OPENSOURCE (from a known open source project)
                        2. THIRDPARTY (commercial third-part SDK)
                        3. FIRSTPARTY
    
                        Please classify this libray into one of these sdk categories:
                        1. Third AI (Provides AI-related machine learning related features)
                        2. Third Crypto (Provides network security functions such as encryption, decryption, signature, signature verification, and SSL)
                        3. Third Compression (Provides data compression and decompression functions)
                        4. Third Audio (Audio capabilities are available)
                        5. Third AVCodec (Provides audio and video encoding and decoding functions)
                        6. Third Media (Provides a media processing framework)
                        7. Third Graphics (Provides graphics-related functions)
                        8. Third Image (Provide image-related functions)
                        9. Third Communication Protocol (Provide network communication protocol related functions)
                        10. Third Data (Provide data storage related functions, including database, serialization, and deserialization)
                        11. Third JSVM (JavaScript virtual machine and runtime)
                        12. Third WEB (WebView related functions)
                        13. Third Cloud (Provide functions related to cloud services)
                        14. Third Map (Map services are available)
                        15. Third DFX (DFX, diagnostics, and logging-related functions are provided)
                        16. Third Test (Testing services are available)
                        17. Third Upload Download (Provides upload and download services functionality)
                        18. Third Startup (Provides startup service functionality)
                        19. Third Payment (Provide payment service features)
                        20. UNKNOWN
    
                        Provide your analysis in this format by json:
                        broad_category: [one of the above categories]
                        specific_origin: [standardized name of the project/company if known]
                        sdk_category：[function speculation ]
                        feature: [function speculation]
                        reasoning: [your explanation]
                        confidence: [number between 0.0 and 1.0] 
                        """
                }]
        }
        response = requests.post(
            "https://api.deepseek.com/v1/chat/completions",
            headers=headers,
            json=payload,
            verify=False
        )

        logger.debug(json.dumps(response.json()))

        if response.status_code == 200:
            content = response.json()['choices'][0]['message']['content']
            content = content[content.find('{'): content.find('}') + 1]
            return json.loads(content)

        return {}

# -------------------- 主流程 --------------------
def analyze_directory(directory):
    analyzer = SOAnalyzer()
    ti = ThreatIntelligence()
    report = {}
    results = []

    for root, _, files in os.walk(directory):
        for file in files:
            if not re.fullmatch(r'^.*\.so(\.\d+)?$', file):
                continue

            try:
                path = os.path.join(root, file)
                logger.info(f"Analyzing {path}...")

                # 基础分析
                result = analyzer.analyze_elf(path)

                deepseek_result = ti.query_deepseek(result)
                logger.info(json.dumps(deepseek_result))
                report[file] = deepseek_result

                result['ti'] = {
                    'deepseek': deepseek_result
                }

                results.append(result)

            except Exception as e:
                logger.error(f"Error analyzing {file}: {str(e)}")

    # 生成报告
    with open('analysis_report.json', 'w') as f:
        json.dump(report, f, indent=2, ensure_ascii=False)

    with open('origin_analysis.json', 'w') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    logger.info(f"生成报告完成，共分析 {len(report)} 个文件")


if __name__ == '__main__':
    target_dir = sys.argv[1]
    analyze_directory(target_dir)


